from paddleocr import PaddleOCR
import cv2
from typing import List, Dict
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor

class OCRProcessor:
    def __init__(self, similarity_threshold: float = 0.8, batch_size: int = 4):
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch")
        self.similarity_threshold = similarity_threshold
        self.batch_size = batch_size
        
    def process(self, image) -> str:
        """处理OCR识别"""
        result = self.ocr.ocr(image, cls=True)
        text = ''
        if result[0]:
            text = ' '.join([line[1][0] for line in result[0]])
        return text
    
    def process_batch(self, image_paths: List[str]) -> List[Dict]:
        """批量处理图片OCR"""
        print("\n开始OCR文本识别...")
        image_info = []
        
        with ThreadPoolExecutor(max_workers=self.batch_size) as executor:
            futures = []
            for img_path in image_paths:
                futures.append(executor.submit(self._process_single_image, img_path))
            
            for future in tqdm(futures, desc="OCR处理进度"):
                result = future.result()
                if result:
                    image_info.append(result)
        
        return image_info
    
    def _process_single_image(self, img_path: str) -> Dict:
        """处理单张图片"""
        img = cv2.imread(img_path)
        if img is not None:
            text = self.process(img)
            return {
                'path': img_path,
                'text': text,
                'quality': self._evaluate_quality(img)
            }
        return None
    
    def deduplicate(self, image_info: List[Dict]) -> List[str]:
        """基于OCR文本内容去重"""
        final_images = []
        used_indices = set()
        
        for i, img1 in enumerate(image_info):
            if i in used_indices:
                continue
                
            similar_group = [i]
            for j in range(i + 1, len(image_info)):
                if j not in used_indices:
                    similarity = self._calculate_text_similarity(
                        img1['text'], 
                        image_info[j]['text']
                    )
                    
                    if similarity > self.similarity_threshold:
                        similar_group.append(j)
                        print(f"\r发现文本相似的图片: {img1['path']} 和 {image_info[j]['path']} "
                              f"(相似度: {similarity:.2f})", end="")
            
            # 在相似组中选择质量最高的图片
            best_idx = max(similar_group, key=lambda x: image_info[x]['quality'])
            final_images.append(image_info[best_idx]['path'])
            used_indices.update(similar_group)
        
        return final_images
    
    def _calculate_text_similarity(self, text1: str, text2: str) -> float:
        """计算两段文本的相似度"""
        if not text1 or not text2:
            return 0.0
            
        # 分词并转换为集合
        words1 = set(text1.split())
        words2 = set(text2.split())
        
        # 计算Jaccard相似度
        intersection = len(words1.intersection(words2))
        union = len(words1.union(words2))
        
        return intersection / union if union > 0 else 0.0
    
    def _evaluate_quality(self, image) -> float:
        """简单的图像质量评估"""
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
            
        # 使用Laplacian方差作为清晰度指标
        laplacian = cv2.Laplacian(gray, cv2.CV_64F)
        return laplacian.var()